import os
import numpy as np
import pandas as pd
from scipy.stats import spearmanr
import glob
from collections import defaultdict

def calculate_spearman_for_files(folder_path):
    # 获取所有符合格式的文件
    file_pattern = os.path.join(folder_path, "stock_*.content")
    files = glob.glob(file_pattern)
    
    # 存储结果的字典
    results = {}
    
    for file_path in files:
        # 从文件名提取日期
        file_name = os.path.basename(file_path)
        date = file_name.replace("stock_", "").replace(".content", "")
        
        # 读取文件内容
        try:
            data = []
            with open(file_path, 'r') as f:
                for line in f:
                    parts = line.strip().split()
                    if len(parts) > 2:  # 确保有足够的列
                        data.append(parts)
            
            if not data:
                print(f"警告: 文件 {file_name} 为空或格式不正确")
                continue
                
            # 转换为numpy数组
            data_array = np.array(data, dtype=object)
            
            # 提取特征和标签
            indices = data_array[:, 0]  # 第一列为索引
            features = data_array[:, 1:-1].astype(float)  # 特征列
            labels = data_array[:, -1].astype(float)  # 最后一列为标签
            
            # 计算所有特征的等权平均
            feature_sum = np.sum(features, axis=1)
            
            # 计算Spearman相关系数
            correlation, p_value = spearmanr(feature_sum, labels)
            
            # 存储结果
            results[date] = {
                'correlation': correlation,
                'p_value': p_value,
                'sample_size': len(data_array)
            }
            
            print(f"文件 {file_name} 的Spearman相关系数: {correlation:.4f} (p-value: {p_value:.4f}, 样本数: {len(data_array)})")
            
        except Exception as e:
            print(f"处理文件 {file_name} 时出错: {str(e)}")
    
    return results

def group_by_year(results):
    # 按年份分组
    year_groups = defaultdict(list)
    
    for date, stats in results.items():
        # 假设日期格式为YYYYMM
        if len(date) >= 4:
            year = date[:4]
            year_groups[year].append(stats['correlation'])
    
    # 计算每年的统计数据
    yearly_stats = {}
    for year, correlations in year_groups.items():
        yearly_stats[year] = {
            'mean': np.mean(correlations),
            'median': np.median(correlations),
            'min': np.min(correlations),
            'max': np.max(correlations),
            'count': len(correlations)
        }
    
    return yearly_stats

def main():
    folder_path = "/root/GAT/data/stock/stock_node_vwap_edge_sec"
    print(f"正在处理文件夹: {folder_path}")
    
    results = calculate_spearman_for_files(folder_path)
    
    # 按年份分组并计算统计数据
    yearly_stats = group_by_year(results)
    
    # 打印年度统计数据
    print("\n按年份统计:")
    print(f"{'年份':<6}{'平均值':<10}{'中位数':<10}{'最小值':<10}{'最大值':<10}{'样本数':<6}")
    print("-" * 52)
    
    # 按年份排序
    for year in sorted(yearly_stats.keys()):
        stats = yearly_stats[year]
        print(f"{year:<6}{stats['mean']:.4f}{'':>5}{stats['median']:.4f}{'':>5}{stats['min']:.4f}{'':>5}{stats['max']:.4f}{'':>5}{stats['count']:<6}")
    
    # 所有年份的总体统计
    all_correlations = [v['correlation'] for v in results.values()]
    print("\n所有数据:")
    print(f"总平均值: {np.mean(all_correlations):.4f}")
    print(f"总中位数: {np.median(all_correlations):.4f}")
    print(f"总样本数: {len(all_correlations)}")

if __name__ == "__main__":
    main()